"""Export LFM2.5-VL-450M as single-method PTE compatible with llama_main."""

from __future__ import annotations

import logging
from pathlib import Path

import torch
from torch.export import Dim
from torch.export._trace import _export
from torch.nn.attention import SDPBackend

from executorch.backends.cuda.cuda_backend import CudaBackend
from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
from executorch.examples.models.lfm2.short_conv import ShortConvBlock
from executorch.examples.models.lfm2_5_vl.model import Lfm2p5VlModel, MAX_SEQ_LEN
from executorch.exir import (
    EdgeCompileConfig,
    ExecutorchBackendConfig,
    to_edge_transform_and_lower,
)
from executorch.exir.passes import MemoryPlanningPass
from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass

logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")

try:
    from torch._inductor.codecache import cuda_compile_utils

    _orig_nvcc_arch = cuda_compile_utils._nvcc_arch_as_compile_option

    def _patched_nvcc_arch() -> str:
        return "103a" if cuda_compile_utils.cuda_env.get_cuda_arch() == "103" else _orig_nvcc_arch()

    cuda_compile_utils._nvcc_arch_as_compile_option = _patched_nvcc_arch
except (ImportError, AttributeError):
    pass

_PARAMS = Path(__file__).parent / ".." / "executorch" / "examples" / "models" / "lfm2_5_vl" / "config" / "lfm2_5_vl_450m_config.json"
_MODEL_DIR = Path("LFM2-VL-450M")
_OUTPUT = Path("lfm2_5_vl_llama_cuda.pte")


class _LlamaCompatModel(torch.nn.Module):
    """forward(input_ids, input_pos) -> logits, matching llama_main interface."""

    def __init__(
        self, lfm2: torch.nn.Module, conv_dim: int, conv_indices: list[int],
        *, dtype: torch.dtype, device: str,
    ) -> None:
        super().__init__()
        self.embed = lfm2.model_.model.language_model.get_input_embeddings()
        self.text_model = lfm2.text_model
        self.conv_indices = conv_indices

        for idx in conv_indices:
            buf = torch.zeros(1, conv_dim, 2, dtype=dtype, device=device)
            self.register_buffer(f"conv_state_{idx}", buf, persistent=False)
            if not torch.compiler.is_compiling():
                torch._dynamo.mark_static_address(buf)

    def forward(self, input_ids: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
        embeddings = self.embed(input_ids)
        conv_states = {idx: getattr(self, f"conv_state_{idx}") for idx in self.conv_indices}
        out = self.text_model(None, {"input_pos": input_pos, "conv_states": conv_states}, embeddings)
        if isinstance(out, tuple):
            out = out[0]
        return out.contiguous()


def main() -> None:
    logging.info("Loading model...")
    lfm2_model = Lfm2p5VlModel(
        model_dir=str(_MODEL_DIR),
        max_seq_len=MAX_SEQ_LEN,
        max_context_len=MAX_SEQ_LEN,
        params_path=str(_PARAMS),
        use_sdpa_with_kv_cache_op=False,
    )
    lfm2 = lfm2_model.get_eager_model().to(dtype=torch.bfloat16, device="cuda")

    conv_indices = [i for i, layer in enumerate(lfm2.text_model.layers) if isinstance(layer, ShortConvBlock)]
    model = _LlamaCompatModel(lfm2, lfm2.text_model_args.dim, conv_indices, dtype=torch.bfloat16, device="cuda")

    # Mark KV cache buffers after device migration
    for module in model.text_model.modules():
        for name, buf in module.named_buffers(recurse=False):
            if name in ("k_cache", "v_cache"):
                torch._dynamo.mark_static_address(buf)

    seq = 8
    token_dim = Dim("token_dim", min=1, max=MAX_SEQ_LEN - 1)
    example_ids = torch.randint(1, 65000, (1, seq), dtype=torch.int64, device="cuda")
    example_pos = torch.arange(seq, dtype=torch.int64, device="cuda")

    logging.info("Exporting...")
    with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad():
        ep = _export(
            model, (example_ids, example_pos),
            dynamic_shapes=({1: token_dim}, {0: token_dim}),
            strict=False,
            prefer_deferred_runtime_asserts_over_guards=True,
        )

    logging.info("Lowering to CUDA")
    compile_specs = [CudaBackend.generate_method_name_compile_spec("forward")]
    et_prog = to_edge_transform_and_lower(
        {"forward": ep},
        partitioner={"forward": [CudaPartitioner(compile_specs)]},
        compile_config=EdgeCompileConfig(_check_ir_validity=False, _skip_dim_order=True),
        constant_methods={"get_max_seq_len": MAX_SEQ_LEN, "get_vocab_size": lfm2.text_model_args.vocab_size, "use_kv_cache": True, "get_eos_ids": [7]},
    )

    et_program = et_prog.to_executorch(
        ExecutorchBackendConfig(
            memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
            sym_shape_eval_pass={"forward": ConstraintBasedSymShapeEvalPass()},
        )
    )

    logging.info("Saving %s", _OUTPUT)
    with open(_OUTPUT, "wb") as f:
        et_program.write_to_file(f)
    et_program.write_tensor_data_to_file(".")
    logging.info("Done")


if __name__ == "__main__":
    main()

"""Run LFM2.5-VL-450M from an exported PTE+PTD on CUDA."""

from __future__ import annotations

import time
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from transformers import AutoProcessor
from executorch.extension.pybindings.portable_lib import _load_for_executorch

PTE_PATH = Path("lfm2_5_vl_bf16_cuda.pte")
PTD_PATH = Path("aoti_cuda_blob.ptd")
MODEL_DIR = Path("LFM2-VL-450M")

IMAGE_TOKEN_ID = 396
EOS_ID = 7
VISION_INPUT_SIZE = 512

# Model card recommended sampling parameters
TEMPERATURE = 0.1
MIN_P = 0.15
REPETITION_PENALTY = 1.05


def _load_image_pixels(path: Path) -> torch.Tensor:
    """Load an image as [1, 3, 512, 512] NCHW float32 in [0, 255]."""
    img = Image.open(path).convert("RGB").resize((VISION_INPUT_SIZE, VISION_INPUT_SIZE))
    return torch.from_numpy(np.array(img)).permute(2, 0, 1).unsqueeze(0).float()


def _embed_tokens(module, input_ids: torch.Tensor) -> torch.Tensor:
    return module.run_method("token_embedding", [input_ids])[0].contiguous()


def _decode_step(module, embeddings: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
    return module.run_method("text_decoder", [embeddings.contiguous(), input_pos])[0]


def _build_embeddings(
    module,
    input_ids: torch.Tensor,
    image_path: Path | None,
) -> torch.Tensor:
    """Build the full embedding sequence, splicing in vision embeddings if needed."""
    if image_path is None or IMAGE_TOKEN_ID not in input_ids[0]:
        return _embed_tokens(module, input_ids)

    positions = (input_ids[0] == IMAGE_TOKEN_ID).nonzero(as_tuple=True)[0]
    first, last = positions[0].item(), positions[-1].item()

    before = _embed_tokens(module, input_ids[:, :first])
    after = _embed_tokens(module, input_ids[:, last + 1 :])

    pixels = _load_image_pixels(image_path).contiguous()
    image = module.run_method("vision_encoder", [pixels])[0].contiguous()

    return torch.cat([before, image, after], dim=1)


def _sample_token(
    logits: torch.Tensor,
    generated: list[int],
    temperature: float = TEMPERATURE,
    min_p: float = MIN_P,
    repetition_penalty: float = REPETITION_PENALTY,
) -> int:
    """Sample next token with temperature, min-p filtering, and repetition penalty."""
    scores = logits.float()

    # Repetition penalty: reduce logits for already-generated tokens
    if generated and repetition_penalty != 1.0:
        prev_tokens = torch.tensor(generated, dtype=torch.long, device=scores.device)
        token_scores = scores[prev_tokens]
        token_scores = torch.where(
            token_scores > 0,
            token_scores / repetition_penalty,
            token_scores * repetition_penalty,
        )
        scores[prev_tokens] = token_scores

    if temperature <= 0:
        return scores.argmax(dim=-1).item()

    probs = torch.softmax(scores / temperature, dim=-1)

    # Min-p filtering: zero out tokens below min_p * max_prob
    if min_p > 0:
        top_prob = probs.max()
        probs[probs < min_p * top_prob] = 0.0

    return torch.multinomial(probs, num_samples=1).item()


def generate(
    module,
    processor: AutoProcessor,
    prompt: str,
    *,
    image_path: Path | None = None,
    max_new_tokens: int = 128,
) -> str:
    content: list[dict[str, str]] = []
    if image_path is not None:
        content.append({"type": "image"})
    content.append({"type": "text", "text": prompt})

    text = processor.apply_chat_template(
        [{"role": "user", "content": content}], add_generation_prompt=True
    )
    input_ids = processor.tokenizer.encode(text, return_tensors="pt")

    embeddings = _build_embeddings(module, input_ids, image_path).contiguous()
    seq_len = embeddings.shape[1]
    logits = _decode_step(module, embeddings, torch.arange(seq_len, dtype=torch.int64))

    generated: list[int] = []
    cur_pos = seq_len

    for _ in range(max_new_tokens):
        last_logits = logits[:, -1, :].squeeze(0) if logits.dim() == 3 else logits.squeeze(0)
        token_id = _sample_token(last_logits, generated)
        if token_id == EOS_ID:
            break
        generated.append(token_id)

        token_embed = _embed_tokens(module, torch.tensor([[token_id]], dtype=torch.int64))
        logits = _decode_step(module, token_embed, torch.tensor([cur_pos], dtype=torch.int64))
        cur_pos += 1

    return processor.tokenizer.decode(generated, skip_special_tokens=True)


def main() -> None:
    module = _load_for_executorch(str(PTE_PATH), str(PTD_PATH))
    processor = AutoProcessor.from_pretrained(str(MODEL_DIR))

    test_image = Path("/tmp/test_image.jpg")

    sections: list[tuple[str, list[tuple[str, Path | None]]]] = [
        (
            "Vision-Language",
            [
                ("Describe this image in detail.", test_image),
                ("What objects do you see in this image?", test_image),
            ],
        ),
        (
            "Text-Only",
            [
                ("The capital of France is", None),
                ("Explain the difference between a compiler and an interpreter in two sentences.", None),
                ("What is the speed of light in meters per second?", None),
            ],
        ),
    ]

    for section_name, prompts in sections:
        print(f"\n=== {section_name} ===")
        for prompt, img in prompts:
            print(f"\nPrompt: {prompt}")
            t0 = time.perf_counter()
            response = generate(module, processor, prompt, image_path=img)
            elapsed = time.perf_counter() - t0
            print(f"Response: {response}")
            print(f"Time: {elapsed:.2f}s")


if __name__ == "__main__":
    main()

"""Benchmark LFM2.5-VL-450M on ExecuTorch CUDA — matches llama_main metrics."""

from __future__ import annotations

import time
from pathlib import Path

import torch
from transformers import AutoProcessor
from executorch.extension.pybindings.portable_lib import _load_for_executorch

PTE_PATH = Path("lfm2_5_vl_bf16_cuda.pte")
PTD_PATH = Path("aoti_cuda_blob.ptd")
MODEL_DIR = Path("LFM2-VL-450M")

EOS_ID = 7
IMAGE_TOKEN_ID = 396

TEMPERATURE = 0.1
TOP_P = 0.9


def _embed(module, ids: torch.Tensor) -> torch.Tensor:
    return module.run_method("token_embedding", [ids])[0].contiguous()


def _decode(module, emb: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
    return module.run_method("text_decoder", [emb.contiguous(), pos])[0]


def _sample(logits: torch.Tensor) -> int:
    if TEMPERATURE <= 0:
        return torch.argmax(logits, dim=-1).item()
    probs = torch.softmax(logits / TEMPERATURE, dim=-1)
    probs_sort, probs_idx = torch.sort(probs, descending=True)
    cum = torch.cumsum(probs_sort, dim=-1)
    mask = (cum - probs_sort) > TOP_P
    probs_sort[mask] = 0.0
    probs_sort /= probs_sort.sum()
    return torch.gather(probs_idx, -1, torch.multinomial(probs_sort, 1)).item()


def benchmark_text(
    module, tokenizer, prompt: str, *, max_new_tokens: int = 128, warmup: bool = True
) -> None:
    tokens = tokenizer.encode(prompt)
    if isinstance(tokens, torch.Tensor):
        tokens = tokens.squeeze().tolist()
    prompt_len = len(tokens)

    if warmup:
        ids = torch.tensor([tokens], dtype=torch.int64)
        emb = _embed(module, ids)
        pos = torch.arange(emb.shape[1], dtype=torch.int64)
        _decode(module, emb, pos)

    # --- Prefill ---
    ids = torch.tensor([tokens], dtype=torch.int64)
    torch.cuda.synchronize()
    t_prefill = time.perf_counter()
    emb = _embed(module, ids)
    pos = torch.arange(emb.shape[1], dtype=torch.int64)
    logits = _decode(module, emb, pos)
    torch.cuda.synchronize()
    prefill_time = time.perf_counter() - t_prefill

    last = logits[:, -1, :].squeeze(0) if logits.dim() == 3 else logits.squeeze(0)
    cur_token = _sample(last)
    generated = [cur_token]
    cur_pos = prompt_len

    # --- Decode ---
    torch.cuda.synchronize()
    t_decode = time.perf_counter()
    while len(generated) < max_new_tokens:
        tok_emb = _embed(module, torch.tensor([[cur_token]], dtype=torch.int64))
        logits = _decode(module, tok_emb, torch.tensor([cur_pos], dtype=torch.int64))
        last = logits[:, -1, :].squeeze(0) if logits.dim() == 3 else logits.squeeze(0)
        cur_token = _sample(last)
        if cur_token == EOS_ID:
            break
        generated.append(cur_token)
        cur_pos += 1
    torch.cuda.synchronize()
    decode_time = time.perf_counter() - t_decode

    decode_tokens = len(generated) - 1  # first token came from prefill
    text = tokenizer.decode(generated, skip_special_tokens=True)

    print(f"  Prompt ({prompt_len} tokens): {prompt[:60]}...")
    print(f"  Output ({len(generated)} tokens): {text[:80]}...")
    print(f"  Prefill:  {prefill_time*1000:.1f} ms  |  {prompt_len/prefill_time:.0f} tok/s")
    print(f"  TTFT:     {prefill_time*1000:.1f} ms")
    if decode_tokens > 0:
        print(f"  Decode:   {decode_time*1000:.1f} ms  |  {decode_tokens/decode_time:.1f} tok/s")
    print()


def benchmark_prefill_sweep(module, tokenizer) -> None:
    """Prefill-only benchmark across different input lengths."""
    print("=== Prefill Sweep ===")
    print(f"{'ISL':>6}  {'Latency (ms)':>12}  {'Throughput (tok/s)':>18}")
    print("-" * 42)

    for isl in [32, 64, 128, 256, 512, 1024]:
        ids = torch.randint(10, 65000, (1, isl), dtype=torch.int64)

        # Warmup
        emb = _embed(module, ids)
        pos = torch.arange(isl, dtype=torch.int64)
        _decode(module, emb, pos)

        # Timed (5 runs, take median)
        times = []
        for _ in range(5):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            emb = _embed(module, ids)
            pos = torch.arange(isl, dtype=torch.int64)
            _decode(module, emb, pos)
            torch.cuda.synchronize()
            times.append(time.perf_counter() - t0)

        times.sort()
        median = times[len(times) // 2]
        print(f"{isl:>6}  {median*1000:>12.2f}  {isl/median:>18.0f}")

    print()


def main() -> None:
    print(f"Loading {PTE_PATH} + {PTD_PATH}\n")
    module = _load_for_executorch(str(PTE_PATH), str(PTD_PATH))
    processor = AutoProcessor.from_pretrained(str(MODEL_DIR))
    tokenizer = processor.tokenizer

    # --- Text generation benchmarks ---
    prompts = [
        "The capital of France is",
        "Explain the difference between a compiler and an interpreter in two sentences.",
        "Write a short paragraph about the history of artificial intelligence.",
        "What is the speed of light in meters per second? Give just the number.",
    ]

    print("=== Text Generation (warmup + timed) ===\n")
    for prompt in prompts:
        benchmark_text(module, tokenizer, prompt, max_new_tokens=64)

    # --- Prefill sweep ---
    benchmark_prefill_sweep(module, tokenizer)


if __name__ == "__main__":
    main()